/*
* Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

use std::str::Utf8Error;

use paste::paste;

pub trait Decoder {
    type Error: Sized;

    fn remaining(&self, buf: &[u8]) -> usize;

    fn get_bool(&mut self, buf: &[u8]) -> Result<bool, Self::Error>;
    fn get_u8(&mut self, buf: &[u8]) -> Result<u8, Self::Error>;
    fn get_u16(&mut self, buf: &[u8]) -> Result<u16, Self::Error>;
    fn get_u32(&mut self, buf: &[u8]) -> Result<u32, Self::Error>;
    fn get_u64(&mut self, buf: &[u8]) -> Result<u64, Self::Error>;
    fn get_i8(&mut self, buf: &[u8]) -> Result<i8, Self::Error>;
    fn get_i16(&mut self, buf: &[u8]) -> Result<i16, Self::Error>;
    fn get_i32(&mut self, buf: &[u8]) -> Result<i32, Self::Error>;
    fn get_i64(&mut self, buf: &[u8]) -> Result<i64, Self::Error>;
    fn get_f32(&mut self, buf: &[u8]) -> Result<f32, Self::Error>;
    fn get_f64(&mut self, buf: &[u8]) -> Result<f64, Self::Error>;
    fn get_bytes<'b>(&mut self, buf: &'b [u8], len: usize) -> Result<&'b [u8], Self::Error>;
    fn get_str<'b>(&mut self, buf: &'b [u8], len: usize) -> Result<&'b str, Self::Error>;

    unsafe fn get_bool_unchecked(&mut self, buf: &[u8]) -> bool;
    unsafe fn get_u8_unchecked(&mut self, buf: &[u8]) -> u8;
    unsafe fn get_u16_unchecked(&mut self, buf: &[u8]) -> u16;
    unsafe fn get_u32_unchecked(&mut self, buf: &[u8]) -> u32;
    unsafe fn get_u64_unchecked(&mut self, buf: &[u8]) -> u64;
    unsafe fn get_i8_unchecked(&mut self, buf: &[u8]) -> i8;
    unsafe fn get_i16_unchecked(&mut self, buf: &[u8]) -> i16;
    unsafe fn get_i32_unchecked(&mut self, buf: &[u8]) -> i32;
    unsafe fn get_i64_unchecked(&mut self, buf: &[u8]) -> i64;
    unsafe fn get_f32_unchecked(&mut self, buf: &[u8]) -> f32;
    unsafe fn get_f64_unchecked(&mut self, buf: &[u8]) -> f64;
    unsafe fn get_bytes_unchecked<'b>(&mut self, buf: &'b [u8], len: usize) -> &'b [u8];
    unsafe fn get_str_unchecked<'b>(&mut self, buf: &'b [u8], len: usize) -> &'b str;
}

#[derive(Default)]
pub struct DefaultDecoder {
    offset: usize,
}

#[derive(Debug, thiserror::Error)]
pub enum DefaultDecoderError {
    #[error("Not enough space in buffer")]
    NotEnoughSpace,
    #[error(transparent)]
    InvalidUtf8(#[from] Utf8Error),
}

macro_rules! impl_get_for_primitive {
    ($type:ident) => {
        paste! {
            fn [<get_ $type>](&mut self, buf: &[u8]) -> Result<$type, Self::Error> {
                if self.offset + size_of::<$type>() > buf.len() {
                    return Err(DefaultDecoderError::NotEnoughSpace);
                }
                Ok(unsafe { self.[<get_ $type _unchecked>](buf) })
            }
        }

        paste! {
            unsafe fn [<get_ $type _unchecked>](&mut self, buf: &[u8]) -> $type {
                let mut bytes = [0u8; size_of::<$type>()];
                bytes.copy_from_slice(&buf[self.offset..self.offset + size_of::<$type>()]);
                let val = $type::from_le_bytes(bytes);
                self.offset += size_of::<$type>();
                val
            }
        }
    };
}

impl Decoder for DefaultDecoder {
    type Error = DefaultDecoderError;

    fn remaining(&self, buf: &[u8]) -> usize {
        buf.len() - self.offset
    }

    fn get_bool(&mut self, buf: &[u8]) -> Result<bool, Self::Error> {
        if self.offset >= buf.len() {
            return Err(DefaultDecoderError::NotEnoughSpace);
        }
        Ok(unsafe { self.get_bool_unchecked(buf) })
    }

    fn get_u8(&mut self, buf: &[u8]) -> Result<u8, Self::Error> {
        if self.offset >= buf.len() {
            return Err(DefaultDecoderError::NotEnoughSpace);
        }
        Ok(unsafe { self.get_u8_unchecked(buf) })
    }

    fn get_i8(&mut self, buf: &[u8]) -> Result<i8, Self::Error> {
        if self.offset >= buf.len() {
            return Err(DefaultDecoderError::NotEnoughSpace);
        }
        Ok(unsafe { self.get_i8_unchecked(buf) })
    }

    fn get_bytes<'b>(&mut self, buf: &'b [u8], len: usize) -> Result<&'b [u8], Self::Error> {
        if self.offset + len > buf.len() {
            return Err(DefaultDecoderError::NotEnoughSpace);
        }
        Ok(unsafe { self.get_bytes_unchecked(buf, len) })
    }

    fn get_str<'b>(&mut self, buf: &'b [u8], len: usize) -> Result<&'b str, Self::Error> {
        let bytes = self.get_bytes(buf, len)?;
        Ok(std::str::from_utf8(bytes)?)
    }

    /// # Examples
    /// 
    /// ```
    /// use ohmw_core::bincode::{Decoder, DefaultDecoder};
    /// 
    /// let buf = [1u8];
    /// let mut decoder = DefaultDecoder::new();
    /// assert_eq!(unsafe { decoder.get_bool_unchecked(&buf) }, true);
    /// ```
    unsafe fn get_bool_unchecked(&mut self, buf: &[u8]) -> bool {
        unsafe { self.get_u8_unchecked(buf) != 0 }
    }

    /// # Examples
    /// 
    /// ```
    /// use ohmw_core::bincode::{Decoder, DefaultDecoder};
    /// 
    /// let buf = [1u8];
    /// let mut decoder = DefaultDecoder::new();
    /// assert_eq!(unsafe { decoder.get_u8_unchecked(&buf) }, 1);
    /// ```
    unsafe fn get_u8_unchecked(&mut self, buf: &[u8]) -> u8 {
        let val = buf[self.offset];
        self.offset += 1;
        val
    }

    unsafe fn get_i8_unchecked(&mut self, buf: &[u8]) -> i8 {
        unsafe { self.get_u8_unchecked(buf) as i8 }
    }
    
    impl_get_for_primitive!(u16);
    impl_get_for_primitive!(u32);
    impl_get_for_primitive!(u64);
    impl_get_for_primitive!(i16);
    impl_get_for_primitive!(i32);
    impl_get_for_primitive!(i64);
    impl_get_for_primitive!(f32);
    impl_get_for_primitive!(f64);

    unsafe fn get_bytes_unchecked<'b>(&mut self, buf: &'b [u8], len: usize) -> &'b [u8] {
        let bytes = &buf[self.offset..self.offset + len];
        self.offset += len;
        bytes
    }

    unsafe fn get_str_unchecked<'b>(&mut self, buf: &'b [u8], len: usize) -> &'b str {
        unsafe { std::str::from_utf8_unchecked(self.get_bytes_unchecked(buf, len)) }
    }
}

impl DefaultDecoder {
    pub fn new() -> Self {
        Self { offset: 0 }
    }
}
